package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"
)

var (
	rootDir          = flag.String("wg_cfg_dir", "wg_all", "The dir of wg config files") // 指定目录
	endpointPrefix   = flag.String("ep_prefix", "Endpoint = ", "The prefix to the endpoint")
	ipPortIdentifier = flag.String("ip_port_id", "IP:PORT", "The identifier to ip:port")
	firstIPCount     = flag.Int("first_ip_count", 3, "The number of random IP addresses to be used")
)

func replaceEndpoint(filePath string, replacement []string) error {
	// 读取文件内容
	content, err := os.ReadFile(filePath)
	if err != nil {
		return err
	}

	ext := filepath.Ext(filePath)
	filePathWithoutExt := filePath[:len(filePath)-len(ext)]

	// 将文件内容按行切割
	lines := strings.Split(string(content), "\n")

	for n, replacestr := range replacement {

		outfilepath := filePathWithoutExt + fmt.Sprintf("_IP%d", n+1) + ".conf"

		// 遍历每一行，查找并替换以"Endpoint = "为开头的行
		for i, line := range lines {
			if strings.HasPrefix(line, *endpointPrefix) {
				lines[i] = *endpointPrefix + replacestr
			}
		}

		// 将修改后的内容拼接为字符串
		newContent := strings.Join(lines, "\n")

		// 写入修改后的内容到文件
		err = os.WriteFile(outfilepath, []byte(newContent), 0644)
		if err != nil {
			return err
		}
	}

	return nil
}

func processFiles(rootDir string, replacement []string) error {
	// 遍历指定目录下的所有文件
	err := filepath.Walk(rootDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// 忽略目录
		if info.IsDir() {
			return nil
		}

		// 只处理.tpl文件
		if strings.HasSuffix(info.Name(), ".tpl") {
			fmt.Printf("Processing file: %s\n", path)

			// 替换文件中的Endpoint
			err := replaceEndpoint(path, replacement)
			if err != nil {
				fmt.Printf("Error processing file %s: %v\n", path, err)
			}
		}

		return nil
	})

	return err
}

func main() {
	flag.Parse()
	// 打开CSV文件
	file, err := os.Open("result.csv")
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	// 创建一个新的Reader
	r := csv.NewReader(file)
	r.ReuseRecord = true

	//查找IP:PORT所在位置
	ip_column := -1
	record, err := r.Read() // read first record
	if err != nil {
		log.Fatal(err)
	}
	for idx, val := range record {
		if val == *ipPortIdentifier {
			ip_column = idx
			break
		}
	}
	if ip_column == -1 {
		log.Fatal("Invalid IP column")
	}

	// 读取CSV文件的第2-3行，第1列的内容
	var data []string
	for i := 1; i <= *firstIPCount; i++ { // 取前3个
		record, err := r.Read()
		if err != nil {
			log.Fatal(err)
		}
		data = append(data, record[ip_column]) // 第1列的内容
	}

	err = processFiles(*rootDir, data)
	if err != nil {
		fmt.Printf("Error processing files: %v\n", err)
	} else {
		fmt.Println("File processing completed successfully.")
	}

	now := time.Now()
	timestamp := now.Unix()

	GenZip(*rootDir, "wg_all_"+fmt.Sprintf("%d", timestamp)+".zip")
}
